{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "依赖"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pandas jsonlines requests tqdm dashscope zhipuai qianfan -i https://mirrors.aliyun.com/pypi/simple/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "生成测试数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re \n",
    "\n",
    "story=''\n",
    "with open('The_Story_of_the_Stone.txt','r',encoding='utf-8') as fp:\n",
    "    for line in fp:\n",
    "        line=line.strip().replace(\"------------\",\" \").replace(\"\\n\", \" \").replace(\" \",\"\")\n",
    "        string_punctuation='!\"#$%&\\'()*+,-./:;<=>?@[\\\\]^_`{|}~'\n",
    "        line=re.sub('[%s]'%string_punctuation,\"\",line)\n",
    "        line=re.sub('[a-zA-Z]','',line)\n",
    "        story+=line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [],
   "source": [
    "reasoning_stars={'32': [16, 116, 43, 70, 59, 106, 8, 48, 112, 67, 25, 101, 82, 93, 76, 62, 6, 18, 108, 4, 34, 53, 85, 90, 126, 22, 45, 121, 39, 96, 75, 30], '64': [16, 116, 43, 70, 59, 106, 8, 48, 112, 67, 25, 101, 82, 93, 76, 62, 6, 18, 108, 4, 34, 53, 85, 90, 126, 22, 45, 121, 39, 96, 75, 30, 154, 168, 256, 197, 140, 200, 184, 229, 195, 172, 210, 174, 163, 131, 238, 214, 222, 159, 177, 242, 236, 189, 216, 228, 146, 184, 248, 207, 244, 151, 136, 143], '128': [16, 116, 43, 70, 59, 106, 8, 48, 112, 67, 25, 101, 82, 93, 76, 62, 6, 18, 108, 4, 34, 53, 85, 90, 126, 22, 45, 121, 39, 96, 75, 30, 154, 168, 256, 197, 140, 200, 184, 229, 195, 172, 210, 174, 163, 131, 238, 214, 222, 159, 177, 242, 236, 189, 216, 228, 146, 184, 248, 207, 244, 151, 136, 143, 332, 472, 306, 510, 417, 462, 370, 359, 374, 438, 381, 262, 343, 290, 364, 430, 478, 422, 484, 492, 394, 406, 385, 441, 293, 271, 272, 487, 500, 345, 445, 302, 286, 466, 392, 324, 378, 473, 403, 265, 410, 412, 501, 328, 493, 453, 398, 284, 339, 312, 318, 354, 300, 458, 309, 426, 328, 280, 448, 506, 433, 360, 256, 352], '256': [16, 116, 43, 70, 59, 106, 8, 48, 112, 67, 25, 101, 82, 93, 76, 62, 6, 18, 108, 4, 34, 53, 85, 90, 126, 22, 45, 121, 39, 96, 75, 30, 154, 168, 256, 197, 140, 200, 184, 229, 195, 172, 210, 174, 163, 131, 238, 214, 222, 159, 177, 242, 236, 189, 216, 228, 146, 184, 248, 207, 244, 151, 136, 143, 332, 472, 306, 510, 417, 462, 370, 359, 374, 438, 381, 262, 343, 290, 364, 430, 478, 422, 484, 492, 394, 406, 385, 441, 293, 271, 272, 487, 500, 345, 445, 302, 286, 466, 392, 324, 378, 473, 403, 265, 410, 412, 501, 328, 493, 453, 398, 284, 339, 312, 318, 354, 300, 458, 309, 426, 328, 280, 448, 506, 433, 360, 256, 352, 533, 638, 724, 729, 605, 795, 652, 700, 1018, 691, 908, 699, 626, 923, 568, 835, 858, 682, 789, 709, 914, 634, 951, 926, 956, 676, 734, 774, 556, 559, 545, 608, 524, 893, 836, 760, 664, 935, 538, 880, 854, 766, 976, 783, 643, 531, 527, 663, 750, 620, 817, 580, 842, 1006, 900, 824, 815, 624, 891, 694, 648, 850, 742, 945, 569, 941, 513, 612, 1015, 754, 1022, 550, 872, 714, 963, 518, 632, 668, 848, 899, 561, 882, 830, 718, 740, 747, 869, 984, 936, 800, 1001, 1010, 990, 590, 956, 659, 886, 725, 577, 760, 998, 988, 685, 995, 966, 779, 599, 979, 808, 786, 908, 704, 806, 587, 574, 593, 544, 863, 675, 771, 828, 602, 803, 866, 644, 928, 918, 971], '512': [16, 116, 43, 70, 59, 106, 8, 48, 112, 67, 25, 101, 82, 93, 76, 62, 6, 18, 108, 4, 34, 53, 85, 90, 126, 22, 45, 121, 39, 96, 75, 30, 154, 168, 256, 197, 140, 200, 184, 229, 195, 172, 210, 174, 163, 131, 238, 214, 222, 159, 177, 242, 236, 189, 216, 228, 146, 184, 248, 207, 244, 151, 136, 143, 332, 472, 306, 510, 417, 462, 370, 359, 374, 438, 381, 262, 343, 290, 364, 430, 478, 422, 484, 492, 394, 406, 385, 441, 293, 271, 272, 487, 500, 345, 445, 302, 286, 466, 392, 324, 378, 473, 403, 265, 410, 412, 501, 328, 493, 453, 398, 284, 339, 312, 318, 354, 300, 458, 309, 426, 328, 280, 448, 506, 433, 360, 256, 352, 533, 638, 724, 729, 605, 795, 652, 700, 1018, 691, 908, 699, 626, 923, 568, 835, 858, 682, 789, 709, 914, 634, 951, 926, 956, 676, 734, 774, 556, 559, 545, 608, 524, 893, 836, 760, 664, 935, 538, 880, 854, 766, 976, 783, 643, 531, 527, 663, 750, 620, 817, 580, 842, 1006, 900, 824, 815, 624, 891, 694, 648, 850, 742, 945, 569, 941, 513, 612, 1015, 754, 1022, 550, 872, 714, 963, 518, 632, 668, 848, 899, 561, 882, 830, 718, 740, 747, 869, 984, 936, 800, 1001, 1010, 990, 590, 956, 659, 886, 725, 577, 760, 998, 988, 685, 995, 966, 779, 599, 979, 808, 786, 908, 704, 806, 587, 574, 593, 544, 863, 675, 771, 828, 602, 803, 866, 644, 928, 918, 971, 1452, 1694, 1665, 1300, 1373, 1651, 1990, 1356, 1150, 1855, 1076, 1616, 2024, 1520, 1467, 1168, 1424, 1876, 1851, 1317, 1534, 1381, 1525, 1522, 1481, 1298, 1962, 1659, 1342, 1440, 1788, 1103, 1162, 1905, 1062, 1717, 1984, 1613, 2038, 1744, 1282, 1746, 1830, 1109, 1166, 1424, 1345, 1184, 1872, 1964, 1812, 1113, 1973, 1271, 2007, 2034, 1572, 1578, 1403, 1608, 1564, 1735, 1754, 1305, 1335, 1601, 1571, 1756, 1952, 1938, 1089, 1393, 2046, 1882, 1934, 1978, 1629, 1380, 1328, 1636, 1472, 1887, 2014, 1230, 1898, 1690, 1627, 1081, 1908, 1588, 1818, 1644, 1685, 1266, 1858, 1415, 1724, 1247, 1474, 1555, 2016, 1292, 1713, 1154, 1860, 1047, 1351, 1774, 1419, 1985, 1360, 1024, 1726, 1644, 1313, 1285, 1030, 1399, 1610, 1730, 1367, 1497, 1997, 1841, 1457, 1486, 1042, 1528, 1145, 1058, 1503, 1386, 1680, 1453, 1506, 1792, 1222, 1739, 1178, 1836, 1309, 1432, 1290, 1957, 1931, 1560, 1182, 1321, 1928, 1632, 1339, 1037, 1444, 1709, 1872, 1768, 1278, 1598, 2000, 1190, 1698, 1128, 1204, 1944, 1804, 1157, 1464, 1194, 1781, 1134, 1827, 1429, 1762, 2010, 1100, 1552, 1772, 1512, 1592, 1077, 1951, 1919, 2030, 1669, 1214, 1332, 1258, 1263, 1118, 1580, 1480, 1088, 1210, 1548, 1900, 1055, 1243, 1495, 1969, 1276, 1893, 1922, 1845, 1409, 1216, 1446, 1707, 1589, 1543, 1510, 2020, 1913, 1354, 1624, 1390, 1678, 1096, 1784, 1172, 1196, 1749, 1814, 1372, 1890, 1703, 1835, 1674, 1778, 1804, 1994, 1406, 1204, 1136, 1237, 1252, 1226, 1660, 1121, 1051, 1034, 1864, 1824, 1948, 1490, 1066, 1798, 1140, 1656, 1234, 1556, 1250, 1126, 1071, 1105, 2042, 1538], '1024': [16, 116, 43, 70, 59, 106, 8, 48, 112, 67, 25, 101, 82, 93, 76, 62, 6, 18, 108, 4, 34, 53, 85, 90, 126, 22, 45, 121, 39, 96, 75, 30, 154, 168, 256, 197, 140, 200, 184, 229, 195, 172, 210, 174, 163, 131, 238, 214, 222, 159, 177, 242, 236, 189, 216, 228, 146, 184, 248, 207, 244, 151, 136, 143, 332, 472, 306, 510, 417, 462, 370, 359, 374, 438, 381, 262, 343, 290, 364, 430, 478, 422, 484, 492, 394, 406, 385, 441, 293, 271, 272, 487, 500, 345, 445, 302, 286, 466, 392, 324, 378, 473, 403, 265, 410, 412, 501, 328, 493, 453, 398, 284, 339, 312, 318, 354, 300, 458, 309, 426, 328, 280, 448, 506, 433, 360, 256, 352, 533, 638, 724, 729, 605, 795, 652, 700, 1018, 691, 908, 699, 626, 923, 568, 835, 858, 682, 789, 709, 914, 634, 951, 926, 956, 676, 734, 774, 556, 559, 545, 608, 524, 893, 836, 760, 664, 935, 538, 880, 854, 766, 976, 783, 643, 531, 527, 663, 750, 620, 817, 580, 842, 1006, 900, 824, 815, 624, 891, 694, 648, 850, 742, 945, 569, 941, 513, 612, 1015, 754, 1022, 550, 872, 714, 963, 518, 632, 668, 848, 899, 561, 882, 830, 718, 740, 747, 869, 984, 936, 800, 1001, 1010, 990, 590, 956, 659, 886, 725, 577, 760, 998, 988, 685, 995, 966, 779, 599, 979, 808, 786, 908, 704, 806, 587, 574, 593, 544, 863, 675, 771, 828, 602, 803, 866, 644, 928, 918, 971, 1452, 1694, 1665, 1300, 1373, 1651, 1990, 1356, 1150, 1855, 1076, 1616, 2024, 1520, 1467, 1168, 1424, 1876, 1851, 1317, 1534, 1381, 1525, 1522, 1481, 1298, 1962, 1659, 1342, 1440, 1788, 1103, 1162, 1905, 1062, 1717, 1984, 1613, 2038, 1744, 1282, 1746, 1830, 1109, 1166, 1424, 1345, 1184, 1872, 1964, 1812, 1113, 1973, 1271, 2007, 2034, 1572, 1578, 1403, 1608, 1564, 1735, 1754, 1305, 1335, 1601, 1571, 1756, 1952, 1938, 1089, 1393, 2046, 1882, 1934, 1978, 1629, 1380, 1328, 1636, 1472, 1887, 2014, 1230, 1898, 1690, 1627, 1081, 1908, 1588, 1818, 1644, 1685, 1266, 1858, 1415, 1724, 1247, 1474, 1555, 2016, 1292, 1713, 1154, 1860, 1047, 1351, 1774, 1419, 1985, 1360, 1024, 1726, 1644, 1313, 1285, 1030, 1399, 1610, 1730, 1367, 1497, 1997, 1841, 1457, 1486, 1042, 1528, 1145, 1058, 1503, 1386, 1680, 1453, 1506, 1792, 1222, 1739, 1178, 1836, 1309, 1432, 1290, 1957, 1931, 1560, 1182, 1321, 1928, 1632, 1339, 1037, 1444, 1709, 1872, 1768, 1278, 1598, 2000, 1190, 1698, 1128, 1204, 1944, 1804, 1157, 1464, 1194, 1781, 1134, 1827, 1429, 1762, 2010, 1100, 1552, 1772, 1512, 1592, 1077, 1951, 1919, 2030, 1669, 1214, 1332, 1258, 1263, 1118, 1580, 1480, 1088, 1210, 1548, 1900, 1055, 1243, 1495, 1969, 1276, 1893, 1922, 1845, 1409, 1216, 1446, 1707, 1589, 1543, 1510, 2020, 1913, 1354, 1624, 1390, 1678, 1096, 1784, 1172, 1196, 1749, 1814, 1372, 1890, 1703, 1835, 1674, 1778, 1804, 1994, 1406, 1204, 1136, 1237, 1252, 1226, 1660, 1121, 1051, 1034, 1864, 1824, 1948, 1490, 1066, 1798, 1140, 1656, 1234, 1556, 1250, 1126, 1071, 1105, 2042, 1538, 3368, 3199, 3803, 2472, 2920, 3584, 3620, 3034, 3195, 4074, 3839, 2599, 2468, 2212, 3104, 2910, 2869, 3833, 2309, 2299, 3499, 3462, 2068, 3643, 2690, 3660, 4060, 3778, 2784, 2092, 3393, 2523, 3963, 2322, 3360, 2354, 2452, 3917, 2485, 3031, 3952, 2901, 3550, 2511, 2155, 3160, 4040, 3474, 2726, 3791, 2761, 2190, 2446, 4012, 3538, 2204, 3524, 3052, 3166, 2706, 2075, 2738, 4018, 2414, 2754, 3402, 2950, 3558, 3357, 3048, 4080, 3450, 2388, 3484, 2960, 3621, 4001, 3665, 2670, 2050, 3601, 2664, 2078, 2579, 2337, 3997, 3591, 3016, 3694, 3699, 3285, 2594, 3564, 2236, 3756, 3648, 2306, 2843, 3060, 3282, 2463, 2164, 3096, 2220, 2908, 3321, 3496, 2545, 2560, 3253, 3325, 3607, 3628, 3346, 3954, 2260, 2820, 3700, 2867, 2278, 4066, 2139, 2408, 3250, 2613, 2232, 3958, 2295, 2610, 3686, 3299, 2376, 3892, 4028, 3732, 3843, 2514, 3862, 2608, 2144, 2796, 2531, 3512, 3376, 3238, 2680, 2861, 3966, 2118, 3712, 2263, 3084, 3413, 2086, 3578, 3234, 3247, 2956, 4071, 3307, 2703, 3912, 3536, 2560, 2151, 3886, 2992, 3750, 3906, 2265, 3762, 2332, 3453, 2364, 2396, 3518, 2721, 3936, 3079, 3680, 2204, 4052, 2243, 3882, 3582, 3086, 2684, 3043, 2502, 2131, 3144, 3478, 3159, 2225, 3704, 3390, 2651, 3095, 2620, 2424, 2159, 2990, 2420, 2378, 2896, 4087, 3506, 2250, 2552, 3639, 2133, 3515, 3736, 2386, 2348, 2120, 3598, 2498, 2093, 2912, 3010, 3338, 2458, 2452, 2274, 2764, 2219, 3026, 2986, 2483, 3112, 2972, 3912, 3993, 4006, 3269, 3168, 3844, 3824, 2178, 3136, 3151, 2658, 3649, 3000, 3832, 2097, 2795, 4013, 2936, 3766, 2810, 2852, 3799, 3399, 2182, 3244, 3273, 2064, 3215, 2965, 3874, 2194, 3671, 3719, 2382, 2709, 3554, 2583, 2126, 2556, 3176, 3120, 3352, 2928, 3291, 3827, 3377, 2934, 2256, 2631, 3212, 2996, 3541, 3661, 2645, 2186, 2341, 3691, 3768, 3386, 3974, 2199, 2812, 2879, 3748, 2634, 3772, 2666, 2570, 3177, 2146, 2772, 4055, 3134, 2674, 3110, 3921, 3979, 3732, 2716, 2782, 3902, 2802, 2591, 3728, 2172, 3406, 2733, 3225, 2830, 2746, 2886, 3434, 3308, 3118, 3154, 4030, 2287, 3430, 3039, 3147, 2348, 2404, 3742, 2832, 3723, 3575, 2980, 3438, 2477, 3990, 2493, 3492, 4038, 3471, 2328, 3344, 2946, 2489, 3889, 2101, 3005, 3933, 3784, 3230, 3854, 2215, 3456, 2396, 3206, 3073, 3220, 2362, 2918, 2742, 3571, 3445, 3418, 3983, 3712, 3187, 2084, 3897, 4078, 3129, 3202, 2638, 2432, 2696, 3222, 2618, 2603, 2964, 3069, 2328, 3443, 2643, 2438, 3527, 2518, 2108, 2698, 2175, 2688, 3546, 2844, 3804, 2821, 3926, 3092, 2859, 3812, 2586, 4088, 2806, 3279, 3382, 3678, 2318, 2358, 3023, 3878, 3755, 3486, 2854, 2436, 2772, 3593, 2536, 3294, 4022, 2892, 2246, 2114, 3850, 3968, 4062, 3266, 3501, 3932, 2164, 2790, 2404, 2235, 2268, 3263, 3793, 2825, 3318, 3530, 2110, 3058, 3302, 2062, 3788, 3102, 3365, 2473, 3675, 3190, 3609, 2839, 2371, 3331, 3626, 2731, 3410, 3332, 2941, 2505, 3810, 2316, 2974, 2928, 2883, 3050, 3864, 3947, 2749, 2534, 3016, 3182, 3316, 3352, 2653, 3656, 3258, 3616, 2574, 2875, 2718, 2054, 3816, 2758, 2896, 2544, 3987, 2627, 2060, 2416, 3872, 3421, 3858, 3065, 4047, 3633, 3424, 2982, 2292, 3467, 2302, 3127, 2283, 4035, 3940, 4093, 2526, 2778, 2442, 2568, 3562]}\n",
    "acquisition_stars={'32': [15, 117, 42, 69, 58, 107, 9, 49, 113, 66, 26, 102, 81, 94, 77, 61, 5, 19, 109, 3, 35, 54, 86, 89, 127, 21, 46, 122, 38, 97, 74, 29], '64': [15, 117, 42, 69, 58, 107, 9, 49, 113, 66, 26, 102, 81, 94, 77, 61, 5, 19, 109, 3, 35, 54, 86, 89, 127, 21, 46, 122, 38, 97, 74, 29, 153, 167, 255, 198, 139, 201, 185, 230, 194, 171, 211, 175, 162, 130, 239, 213, 223, 158, 178, 243, 235, 190, 217, 227, 145, 183, 249, 206, 245, 150, 135, 142], '128': [15, 117, 42, 69, 58, 107, 9, 49, 113, 66, 26, 102, 81, 94, 77, 61, 5, 19, 109, 3, 35, 54, 86, 89, 127, 21, 46, 122, 38, 97, 74, 29, 153, 167, 255, 198, 139, 201, 185, 230, 194, 171, 211, 175, 162, 130, 239, 213, 223, 158, 178, 243, 235, 190, 217, 227, 145, 183, 249, 206, 245, 150, 135, 142, 333, 471, 305, 509, 418, 461, 371, 358, 373, 437, 382, 263, 342, 291, 365, 429, 479, 421, 483, 491, 393, 407, 386, 442, 294, 270, 273, 486, 499, 346, 446, 301, 287, 467, 391, 323, 379, 474, 402, 266, 409, 413, 502, 329, 494, 454, 397, 283, 338, 313, 319, 353, 299, 457, 310, 427, 327, 279, 449, 505, 434, 361, 257, 351], '256': [15, 117, 42, 69, 58, 107, 9, 49, 113, 66, 26, 102, 81, 94, 77, 61, 5, 19, 109, 3, 35, 54, 86, 89, 127, 21, 46, 122, 38, 97, 74, 29, 153, 167, 255, 198, 139, 201, 185, 230, 194, 171, 211, 175, 162, 130, 239, 213, 223, 158, 178, 243, 235, 190, 217, 227, 145, 183, 249, 206, 245, 150, 135, 142, 333, 471, 305, 509, 418, 461, 371, 358, 373, 437, 382, 263, 342, 291, 365, 429, 479, 421, 483, 491, 393, 407, 386, 442, 294, 270, 273, 486, 499, 346, 446, 301, 287, 467, 391, 323, 379, 474, 402, 266, 409, 413, 502, 329, 494, 454, 397, 283, 338, 313, 319, 353, 299, 457, 310, 427, 327, 279, 449, 505, 434, 361, 257, 351, 534, 639, 723, 730, 606, 794, 653, 701, 1019, 690, 909, 698, 625, 922, 567, 834, 859, 683, 790, 710, 913, 633, 950, 927, 957, 677, 735, 773, 555, 558, 546, 609, 523, 894, 837, 759, 665, 934, 539, 879, 853, 765, 975, 782, 642, 530, 526, 662, 749, 619, 818, 581, 843, 1007, 901, 823, 814, 623, 890, 693, 649, 851, 741, 946, 570, 942, 514, 613, 1014, 755, 1021, 551, 873, 715, 962, 517, 631, 669, 847, 898, 562, 883, 829, 717, 739, 746, 870, 983, 937, 799, 1002, 1011, 991, 589, 955, 658, 885, 726, 578, 761, 997, 987, 686, 994, 967, 778, 598, 978, 809, 785, 907, 705, 807, 586, 575, 594, 543, 862, 674, 770, 827, 601, 802, 867, 645, 929, 919, 970], '512': [15, 117, 42, 69, 58, 107, 9, 49, 113, 66, 26, 102, 81, 94, 77, 61, 5, 19, 109, 3, 35, 54, 86, 89, 127, 21, 46, 122, 38, 97, 74, 29, 153, 167, 255, 198, 139, 201, 185, 230, 194, 171, 211, 175, 162, 130, 239, 213, 223, 158, 178, 243, 235, 190, 217, 227, 145, 183, 249, 206, 245, 150, 135, 142, 333, 471, 305, 509, 418, 461, 371, 358, 373, 437, 382, 263, 342, 291, 365, 429, 479, 421, 483, 491, 393, 407, 386, 442, 294, 270, 273, 486, 499, 346, 446, 301, 287, 467, 391, 323, 379, 474, 402, 266, 409, 413, 502, 329, 494, 454, 397, 283, 338, 313, 319, 353, 299, 457, 310, 427, 327, 279, 449, 505, 434, 361, 257, 351, 534, 639, 723, 730, 606, 794, 653, 701, 1019, 690, 909, 698, 625, 922, 567, 834, 859, 683, 790, 710, 913, 633, 950, 927, 957, 677, 735, 773, 555, 558, 546, 609, 523, 894, 837, 759, 665, 934, 539, 879, 853, 765, 975, 782, 642, 530, 526, 662, 749, 619, 818, 581, 843, 1007, 901, 823, 814, 623, 890, 693, 649, 851, 741, 946, 570, 942, 514, 613, 1014, 755, 1021, 551, 873, 715, 962, 517, 631, 669, 847, 898, 562, 883, 829, 717, 739, 746, 870, 983, 937, 799, 1002, 1011, 991, 589, 955, 658, 885, 726, 578, 761, 997, 987, 686, 994, 967, 778, 598, 978, 809, 785, 907, 705, 807, 586, 575, 594, 543, 862, 674, 770, 827, 601, 802, 867, 645, 929, 919, 970, 1451, 1695, 1666, 1301, 1374, 1650, 1991, 1357, 1149, 1854, 1075, 1617, 2025, 1519, 1466, 1169, 1423, 1877, 1850, 1318, 1533, 1382, 1526, 1523, 1482, 1297, 1961, 1658, 1341, 1439, 1789, 1102, 1163, 1906, 1061, 1718, 1983, 1614, 2037, 1743, 1281, 1747, 1831, 1110, 1165, 1425, 1346, 1185, 1871, 1965, 1811, 1114, 1974, 1270, 2006, 2033, 1573, 1579, 1402, 1607, 1565, 1734, 1755, 1306, 1334, 1602, 1570, 1757, 1953, 1937, 1090, 1394, 2045, 1883, 1933, 1979, 1630, 1379, 1327, 1637, 1471, 1886, 2015, 1231, 1899, 1691, 1626, 1082, 1909, 1587, 1819, 1645, 1686, 1267, 1857, 1414, 1723, 1246, 1473, 1554, 2017, 1293, 1714, 1155, 1861, 1046, 1350, 1773, 1418, 1986, 1361, 1025, 1727, 1643, 1314, 1286, 1031, 1398, 1609, 1729, 1366, 1498, 1998, 1842, 1458, 1485, 1041, 1529, 1146, 1059, 1502, 1387, 1681, 1454, 1507, 1793, 1223, 1738, 1177, 1837, 1310, 1433, 1289, 1958, 1930, 1561, 1181, 1322, 1927, 1633, 1338, 1038, 1443, 1710, 1873, 1767, 1277, 1599, 2001, 1189, 1699, 1129, 1203, 1943, 1805, 1158, 1463, 1195, 1782, 1133, 1826, 1430, 1763, 2011, 1099, 1551, 1771, 1513, 1593, 1078, 1950, 1918, 2029, 1670, 1213, 1331, 1257, 1262, 1119, 1581, 1479, 1087, 1209, 1547, 1901, 1054, 1242, 1494, 1970, 1275, 1894, 1923, 1846, 1410, 1217, 1447, 1706, 1590, 1542, 1511, 2021, 1914, 1355, 1623, 1391, 1677, 1095, 1785, 1173, 1197, 1750, 1815, 1371, 1889, 1702, 1834, 1675, 1779, 1803, 1995, 1405, 1205, 1137, 1238, 1253, 1227, 1661, 1122, 1050, 1033, 1865, 1823, 1947, 1491, 1067, 1799, 1141, 1655, 1235, 1557, 1251, 1125, 1070, 1106, 2041, 1537], '1024': [15, 117, 42, 69, 58, 107, 9, 49, 113, 66, 26, 102, 81, 94, 77, 61, 5, 19, 109, 3, 35, 54, 86, 89, 127, 21, 46, 122, 38, 97, 74, 29, 153, 167, 255, 198, 139, 201, 185, 230, 194, 171, 211, 175, 162, 130, 239, 213, 223, 158, 178, 243, 235, 190, 217, 227, 145, 183, 249, 206, 245, 150, 135, 142, 333, 471, 305, 509, 418, 461, 371, 358, 373, 437, 382, 263, 342, 291, 365, 429, 479, 421, 483, 491, 393, 407, 386, 442, 294, 270, 273, 486, 499, 346, 446, 301, 287, 467, 391, 323, 379, 474, 402, 266, 409, 413, 502, 329, 494, 454, 397, 283, 338, 313, 319, 353, 299, 457, 310, 427, 327, 279, 449, 505, 434, 361, 257, 351, 534, 639, 723, 730, 606, 794, 653, 701, 1019, 690, 909, 698, 625, 922, 567, 834, 859, 683, 790, 710, 913, 633, 950, 927, 957, 677, 735, 773, 555, 558, 546, 609, 523, 894, 837, 759, 665, 934, 539, 879, 853, 765, 975, 782, 642, 530, 526, 662, 749, 619, 818, 581, 843, 1007, 901, 823, 814, 623, 890, 693, 649, 851, 741, 946, 570, 942, 514, 613, 1014, 755, 1021, 551, 873, 715, 962, 517, 631, 669, 847, 898, 562, 883, 829, 717, 739, 746, 870, 983, 937, 799, 1002, 1011, 991, 589, 955, 658, 885, 726, 578, 761, 997, 987, 686, 994, 967, 778, 598, 978, 809, 785, 907, 705, 807, 586, 575, 594, 543, 862, 674, 770, 827, 601, 802, 867, 645, 929, 919, 970, 1451, 1695, 1666, 1301, 1374, 1650, 1991, 1357, 1149, 1854, 1075, 1617, 2025, 1519, 1466, 1169, 1423, 1877, 1850, 1318, 1533, 1382, 1526, 1523, 1482, 1297, 1961, 1658, 1341, 1439, 1789, 1102, 1163, 1906, 1061, 1718, 1983, 1614, 2037, 1743, 1281, 1747, 1831, 1110, 1165, 1425, 1346, 1185, 1871, 1965, 1811, 1114, 1974, 1270, 2006, 2033, 1573, 1579, 1402, 1607, 1565, 1734, 1755, 1306, 1334, 1602, 1570, 1757, 1953, 1937, 1090, 1394, 2045, 1883, 1933, 1979, 1630, 1379, 1327, 1637, 1471, 1886, 2015, 1231, 1899, 1691, 1626, 1082, 1909, 1587, 1819, 1645, 1686, 1267, 1857, 1414, 1723, 1246, 1473, 1554, 2017, 1293, 1714, 1155, 1861, 1046, 1350, 1773, 1418, 1986, 1361, 1025, 1727, 1643, 1314, 1286, 1031, 1398, 1609, 1729, 1366, 1498, 1998, 1842, 1458, 1485, 1041, 1529, 1146, 1059, 1502, 1387, 1681, 1454, 1507, 1793, 1223, 1738, 1177, 1837, 1310, 1433, 1289, 1958, 1930, 1561, 1181, 1322, 1927, 1633, 1338, 1038, 1443, 1710, 1873, 1767, 1277, 1599, 2001, 1189, 1699, 1129, 1203, 1943, 1805, 1158, 1463, 1195, 1782, 1133, 1826, 1430, 1763, 2011, 1099, 1551, 1771, 1513, 1593, 1078, 1950, 1918, 2029, 1670, 1213, 1331, 1257, 1262, 1119, 1581, 1479, 1087, 1209, 1547, 1901, 1054, 1242, 1494, 1970, 1275, 1894, 1923, 1846, 1410, 1217, 1447, 1706, 1590, 1542, 1511, 2021, 1914, 1355, 1623, 1391, 1677, 1095, 1785, 1173, 1197, 1750, 1815, 1371, 1889, 1702, 1834, 1675, 1779, 1803, 1995, 1405, 1205, 1137, 1238, 1253, 1227, 1661, 1122, 1050, 1033, 1865, 1823, 1947, 1491, 1067, 1799, 1141, 1655, 1235, 1557, 1251, 1125, 1070, 1106, 2041, 1537, 3369, 3198, 3802, 2471, 2921, 3585, 3619, 3035, 3194, 4075, 3838, 2598, 2467, 2211, 3105, 2909, 2870, 3834, 2310, 2298, 3498, 3461, 2069, 3642, 2691, 3659, 4059, 3779, 2785, 2091, 3394, 2522, 3962, 2323, 3361, 2353, 2451, 3918, 2486, 3030, 3951, 2902, 3549, 2510, 2154, 3161, 4041, 3475, 2727, 3790, 2762, 2191, 2447, 4011, 3537, 2203, 3523, 3053, 3165, 2705, 2074, 2739, 4019, 2413, 2755, 3403, 2949, 3559, 3358, 3047, 4081, 3449, 2389, 3483, 2959, 3622, 4002, 3666, 2669, 2049, 3602, 2663, 2077, 2578, 2338, 3998, 3590, 3017, 3693, 3698, 3286, 2593, 3565, 2237, 3757, 3647, 2307, 2842, 3061, 3281, 2462, 2165, 3097, 2221, 2907, 3322, 3495, 2546, 2559, 3254, 3326, 3606, 3629, 3345, 3953, 2259, 2819, 3701, 2866, 2277, 4067, 2138, 2409, 3249, 2614, 2231, 3957, 2294, 2611, 3685, 3298, 2375, 3893, 4027, 3733, 3842, 2513, 3861, 2607, 2143, 2797, 2530, 3511, 3375, 3237, 2679, 2862, 3967, 2119, 3713, 2262, 3083, 3414, 2085, 3579, 3235, 3246, 2955, 4070, 3306, 2702, 3913, 3535, 2561, 2150, 3885, 2993, 3751, 3905, 2266, 3761, 2333, 3454, 2365, 2395, 3517, 2722, 3937, 3078, 3681, 2205, 4051, 2242, 3881, 3581, 3087, 2683, 3042, 2503, 2130, 3143, 3477, 3158, 2226, 3705, 3389, 2650, 3094, 2621, 2425, 2158, 2991, 2421, 2377, 2897, 4086, 3505, 2251, 2551, 3638, 2134, 3514, 3737, 2385, 2349, 2121, 3597, 2499, 2094, 2913, 3009, 3339, 2457, 2453, 2273, 2765, 2218, 3027, 2987, 2482, 3113, 2971, 3911, 3994, 4007, 3270, 3169, 3845, 3823, 2179, 3137, 3150, 2657, 3650, 3001, 3831, 2098, 2794, 4014, 2937, 3765, 2811, 2851, 3798, 3398, 2183, 3243, 3274, 2065, 3214, 2966, 3873, 2195, 3670, 3718, 2383, 2710, 3555, 2582, 2127, 2555, 3175, 3121, 3351, 2927, 3290, 3826, 3378, 2935, 2255, 2630, 3211, 2997, 3542, 3662, 2646, 2185, 2342, 3690, 3769, 3387, 3975, 2198, 2813, 2878, 3747, 2635, 3773, 2667, 2569, 3178, 2145, 2773, 4054, 3135, 2675, 3111, 3922, 3978, 3731, 2715, 2781, 3901, 2803, 2590, 3727, 2171, 3405, 2734, 3226, 2831, 2745, 2887, 3435, 3309, 3117, 3153, 4029, 2286, 3431, 3038, 3146, 2347, 2405, 3741, 2833, 3722, 3574, 2979, 3437, 2478, 3989, 2494, 3491, 4037, 3470, 2329, 3343, 2945, 2490, 3890, 2102, 3006, 3934, 3783, 3229, 3853, 2214, 3457, 2397, 3205, 3074, 3219, 2361, 2919, 2741, 3570, 3446, 3417, 3982, 3711, 3186, 2083, 3898, 4079, 3130, 3201, 2637, 2431, 2695, 3221, 2617, 2602, 2963, 3070, 2327, 3442, 2642, 2437, 3526, 2517, 2107, 2699, 2174, 2687, 3545, 2845, 3805, 2822, 3927, 3091, 2858, 3813, 2585, 4089, 2805, 3278, 3381, 3679, 2319, 2359, 3022, 3877, 3754, 3485, 2853, 2435, 2771, 3594, 2537, 3293, 4023, 2891, 2247, 2113, 3849, 3969, 4063, 3265, 3502, 3931, 2163, 2789, 2403, 2234, 2269, 3262, 3794, 2826, 3317, 3531, 2109, 3059, 3303, 2063, 3787, 3103, 3366, 2474, 3674, 3191, 3610, 2838, 2370, 3330, 3625, 2730, 3409, 3333, 2942, 2506, 3811, 2315, 2973, 2929, 2882, 3051, 3865, 3946, 2750, 2535, 3015, 3181, 3315, 3353, 2654, 3655, 3259, 3615, 2573, 2874, 2719, 2055, 3817, 2757, 2895, 2543, 3986, 2626, 2059, 2417, 3871, 3422, 3859, 3066, 4046, 3634, 3425, 2981, 2291, 3466, 2301, 3126, 2282, 4034, 3941, 4094, 2527, 2779, 2441, 2567, 3561]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [],
   "source": [
    "acquisition_question = \"\\n\\n\\n\\n在这个月光皎洁、云雾缭绕的夜晚，小企鹅正望向天空，全神贯注地数★。请帮助小企鹅收集所数★的颗数，按照如下格式：{\\\"小企鹅\\\":[x,x,x,...]}，不要求和，[x,x,x,...]中数字为小企鹅每次数★的颗数，仅以JSON格式输出结果，不需要输出任何解释。\"\n",
    "reasoning_question = \"\\n\\n\\n\\n在这个月光皎洁、云雾缭绕的夜晚，小企鹅正望向天空，全神贯注地数★。请帮助小企鹅收集所数★的正确颗数，按照如下格式：{\\\"小企鹅\\\":[x,x,x,...]}，不要求和，[x,x,x,...]中数字为小企鹅正确数★的颗数，仅以JSON格式输出结果，不需要输出任何解释。\"        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 生成sample_count个样本，长度按step_size递增，每个sample插入num_needles根针\n",
    "def generate_test_case(sample_count,max_seq_len,num_needles,test_type):\n",
    "    ret=[]\n",
    "    \n",
    "    step_len=int(max_seq_len/sample_count)\n",
    "    for i in range(sample_count):        \n",
    "        # 生成needles\n",
    "        if test_type=='acquisition':\n",
    "            needles=[f'\\n小企鹅数了{r}颗★\\n' for r in acquisition_stars[str(num_needles)]]\n",
    "        else:\n",
    "            needles=[f'\\n小企鹅数了{a}颗★，但发现数错了，于是又数了一遍，这次数对了，是{r}颗★\\n' for (a,r) in zip(reasoning_stars[str(num_needles)],acquisition_stars[str(num_needles)])]\n",
    "        \n",
    "        # 原文\n",
    "        sample_size=int(step_len*(i+1)*0.7) # 让出30%字数给needles\n",
    "        sample_content=story[:sample_size]\n",
    "        \n",
    "        # 插入needles\n",
    "        parts=[]\n",
    "        needle_step=int(sample_size/num_needles)\n",
    "        for j in range(num_needles):\n",
    "            parts.append(sample_content[j*needle_step:(j+1)*needle_step])\n",
    "            parts.append(needles[j])\n",
    "        sample_content=''.join(parts)\n",
    "        \n",
    "        # 生成prompt\n",
    "        prompt=sample_content+(acquisition_question if test_type=='acquisition' else reasoning_question)\n",
    "        ret.append({'prompt':prompt,'test_type':test_type,'sample_size':len(sample_content),'num_needles':num_needles,'acquisition_stars':acquisition_stars[str(num_needles)],'reasoning_stars':reasoning_stars[str(num_needles)]})\n",
    "    return ret "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dashscope\n",
    "dashscope.api_key=''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通义千文"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from http import HTTPStatus\n",
    "from dashscope import Generation  # 建议dashscope SDK 的版本 >= 1.14.0\n",
    "\n",
    "def qwen_call(model,prompt):\n",
    "    messages = [{'role': 'system', 'content': 'You are a helpful assistant.'},\n",
    "                {'role': 'user', 'content': prompt}]\n",
    "    response = Generation.call(model=model,\n",
    "                               messages=messages,\n",
    "                               # 设置随机数种子seed，如果没有设置，则随机数种子默认为1234\n",
    "                               seed=random.randint(1, 10000),\n",
    "                               # 将输出设置为\"message\"格式\n",
    "                               result_format='message')\n",
    "\n",
    "    return response.output['choices'][0]['message']['content']\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "智谱GLM4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zhipuai import ZhipuAI\n",
    "zhipu=ZhipuAI(api_key=\"\") # 填写您自己的APIKey"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zhipu_call(model,prompt):\n",
    "    response=zhipu.chat.completions.create(\n",
    "        model=model,  # 填写需要调用的模型名称\n",
    "        messages=[\n",
    "            {\"role\": \"user\", \"content\":prompt},\n",
    "        ],\n",
    "    )\n",
    "    return response.choices[0].message.content\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "百度文心"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import qianfan\n",
    "\n",
    "# 使用安全认证AK/SK鉴权，通过环境变量初始化；替换下列示例中参数，安全认证Access Key替换your_iam_ak，Secret Key替换your_iam_sk\n",
    "os.environ[\"QIANFAN_ACCESS_KEY\"] = \"\"\n",
    "os.environ[\"QIANFAN_SECRET_KEY\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {},
   "outputs": [],
   "source": [
    "def wenxin_call(model,prompt):\n",
    "    chat_comp = qianfan.ChatCompletion()\n",
    "    # 指定特定模型\n",
    "    resp = chat_comp.do(model=model, messages=[{\n",
    "        \"role\": \"user\",\n",
    "        \"content\": prompt\n",
    "    }])\n",
    "    return resp['body']['result']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json \n",
    "# 评测函数\n",
    "def run_test_case(model_fn,model_name,test_cases):\n",
    "    stats=[]\n",
    "    for sample in test_cases:\n",
    "        resp=model_fn(model_name,sample['prompt'])\n",
    "        resp=resp.strip('`')\n",
    "        if resp.startswith('json'):\n",
    "            resp=resp[4:]\n",
    "        pred_stars=json.loads(resp)\n",
    "        pred_stars=set(pred_stars['小企鹅'])\n",
    "\n",
    "        \n",
    "        if sample['test_type']=='acquisition':\n",
    "            for a in sample['acquisition_stars']:\n",
    "                if a in pred_stars:\n",
    "                    score=1\n",
    "                else:\n",
    "                    score=0\n",
    "                stats.append(score)\n",
    "        else:\n",
    "            for a,r in zip(sample['acquisition_stars'],sample['reasoning_stars']):\n",
    "                if a in pred_stars and r not in pred_stars:\n",
    "                    score=1\n",
    "                elif a in pred_stars and r in pred_stars:\n",
    "                    score=0.5\n",
    "                elif a not in pred_stars and r in pred_stars:\n",
    "                    score=0.25\n",
    "                else:\n",
    "                    score=0\n",
    "                stats.append(score)\n",
    "    print('acc:%.2f'%(sum(stats)/len(stats)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 生成测试集\n",
    "test_cases=generate_test_case(32,6000,32,'reasoning')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc:0.92\n"
     ]
    }
   ],
   "source": [
    "# qwen\n",
    "run_test_case(qwen_call,'qwen-plus',test_cases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc:0.96\n"
     ]
    }
   ],
   "source": [
    "# zhipu\n",
    "run_test_case(zhipu_call,'glm-4',test_cases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc:1.00\n"
     ]
    }
   ],
   "source": [
    "# wenxin\n",
    "run_test_case(wenxin_call,'ERNIE-4.0-8K-0329',test_cases)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "countingstars",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
